import torch
import torchvision
import torchvision.transforms as transforms

from ANN1 import SNN

snn = SNN(4, (784, 500, 500, 10))
snn.load()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

arr = []
for t in range(16,17):
    print(f"timeStep = {t}")
    train_dataset = torchvision.datasets.MNIST(root="data", train=True, download=False, transform=transforms.ToTensor())
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)


    n = 0
    L = 0
    for i, (image, label) in enumerate(train_loader):
        if (int(i) - 1) % 1000 == 0:
            print(f"{n}/{int(i)}:{n/int(i)}")
        if (int(i) + 1) % 10000 == 0:
            arr.append([t,n/int(i)])
            break
        if snn.run(image,t).argmax() == int(label):
            n += 1

for i in arr:
    print(i[0], end=' ')
    print(i[1])